package com.test.daily.leetcode.y2022.m06.day0625.v04;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * @author Tom on 2022/6/25
 */
public class Solution {


    public  class UnionFind {
        // parent[i] = k ： i的父亲是k
        private int[] parent;
        // size[i] = k ： 如果i是代表节点，size[i]才有意义，否则无意义
        // i所在的集合大小是多少
        private int[] size;
        // 辅助结构
        private int[] help;
        // 一共有多少个集合
        private int sets;

        public UnionFind(int N) {
            parent = new int[N];
            size = new int[N];
            help = new int[N];
            sets = N;
            for (int i = 0; i < N; i++) {
                parent[i] = i;
                size[i] = 1;
            }
        }

        // 从i开始一直往上，往上到不能再往上，代表节点，返回
        // 这个过程要做路径压缩
        private int find(int i) {
            int hi = 0;
            while (i != parent[i]) {
                help[hi++] = i;
                i = parent[i];
            }
            for (hi--; hi >= 0; hi--) {
                parent[help[hi]] = i;
            }
            return i;
        }

        public void union(int i, int j) {
            int f1 = find(i);
            int f2 = find(j);
            if (f1 != f2) {
                if (size[f1] >= size[f2]) {
                    size[f1] += size[f2];
                    parent[f2] = f1;
                } else {
                    size[f2] += size[f1];
                    parent[f1] = f2;
                }
                sets--;
            }
        }
        public int sets() {
            return sets;
        }
    }
    public long countPairs(int n, int[][] edges) {
        UnionFind uf = new UnionFind(edges.length);
        for(int i = 0; i < edges.length; i++){
            uf.union(edges[i][0], edges[i][1]);
        }
        HashMap<Integer, List<Integer>> map = new HashMap<>();
        for(int i = 0; i < edges.length; i++){
            if(map.get(uf.find(edges[i][0])) == null){
                List<Integer> list = new ArrayList<>();
                list.add(edges[i][0]);
                map.put(edges[i][0], list);
            } else {
                map.get(uf.find(edges[i][0])).add(edges[i][0]);
                map.get(uf.find(edges[i][0])).add(edges[i][1]);
            }
        }
        int connectedCount = 0;
        for (List<Integer> list: map.values()){
            connectedCount += list.size()*(list.size()-1)/2;
        }
        return n * (n-1) / 2 - connectedCount;
    }
}
